import pandas as pd
import statsmodels.formula.api as sm

from cat_rfs.code.utils.table import Table

SETTINGS = {'main': {'beta_var': 'beta_sheet',
                     'yield_var': 'sheet_yield_usd',
                     'table_file': 'table5'},
            'trace': {'beta_var': 'beta_act',
                      'yield_var': 'act_yield_usd',
                      'table_file': 'table5_trace'},
            'loss': {'beta_var': 'beta_loss',
                     'yield_var': 'sheet_yield_usd',
                     'table_file': 'table5_loss'},
            }


def print_table5(df, output='console', setting='main'):
    beta = SETTINGS[setting]['beta_var']
    yield_ = SETTINGS[setting]['yield_var']
    tname = SETTINGS[setting]['table_file']

    df = df.copy()
    df = df.rename(columns={yield_: 'yield'})
    df = df.sort_values(['date', 'CUSIP9']).reset_index(drop=True).copy()

    df['er'] = df['yield'] - df['expected_loss'] / 100 - df['rf'] / 100
    df['CUSIP6'] = df['CUSIP9'].str[:6]

    df = df.dropna(subset=[beta, 'er'])

    # %% Calculate peril weights
    perils = list(df.loc[~df['perils'].str.contains(','), 'perils'].unique())
    perils = ['el_' + s.replace(' ', '_').lower() for s in perils]
    df['totel'] = df[perils].sum(axis='columns')

    for n, peril in enumerate(perils):
        df[f'{peril}_weight'] = df[peril] / df['totel']
        df[f'{peril}_amt'] = df[f'{peril}_weight'] * df['size']
        df[f'{peril}_amt_group'] = df.groupby(['date'])[f'{peril}_amt'].transform('sum')
        df[f'{peril}_weighted'] = df[f'{peril}_weight'] * df[f'{peril}_amt_group']

    df['weight_peril'] = df[[f'{peril}_weighted' for peril in perils]].sum(axis='columns') / df[
        [f'{peril}_amt_group' for peril in perils]].sum(axis='columns')
    # %% End

    df['er_m'] = df['er'] * df['size']
    df2 = df.groupby('date', as_index=False)[['er_m', 'size']].sum()
    df2['er_m'] = df2['er_m'] / df2['size']

    df = df.drop(columns=['er_m'])
    df = df.merge(df2[['date', 'er_m']], on='date')

    # %% Get Size/AUM
    df2 = pd.read_csv('cat_rfs/data/ts.csv', parse_dates=['date'])
    df2['lev'] = df2['alternative_capital'] / df2['aum']

    df = df.merge(df2, on='date')

    # %% Initialize estimate lists
    const = ['$F_t$']
    const_se = ['']

    el_coef = ['$el_i F_t$']
    el_coef_se = ['']

    w_coef = ['$w_{c,t} F_t$']
    w_coef_se = ['']

    R2 = ['$R^2$']
    N = ['$N$']

    # %% Beta_sheet
    df['F'] = df['var_m_sheet'] * df['lev']
    ols = sm.ols(f'er ~ -1 + F : {beta}', data=df).fit(
        cov_type='cluster', cov_kwds={'groups': df['CUSIP6']}, use_t=True)
    const.append(ols.params[0])
    const_se.append(ols.bse[0])
    el_coef.append('')
    el_coef_se.append('')
    w_coef.append('')
    w_coef_se.append('')
    R2 += [ols.rsquared]
    N += [ols.nobs]

    df['F'] = df['lev']
    ols = sm.ols(f'er ~ -1 + F : {beta}', data=df).fit(
        cov_type='cluster', cov_kwds={'groups': df['CUSIP6']}, use_t=True)
    const.append(ols.params[0])
    const_se.append(ols.bse[0])
    el_coef.append('')
    el_coef_se.append('')
    w_coef.append('')
    w_coef_se.append('')
    R2 += [ols.rsquared]
    N += [ols.nobs]

    df['F'] = df['er_m']
    ols = sm.ols(f'er ~ -1 + F : {beta}', data=df).fit(
        cov_type='cluster', cov_kwds={'groups': df['CUSIP6']}, use_t=True)
    const.append(ols.params[0])
    const_se.append(ols.bse[0])
    el_coef.append('')
    el_coef_se.append('')
    w_coef.append('')
    w_coef_se.append('')
    R2 += [ols.rsquared]
    N += [ols.nobs]

    # %% Characteristics
    df['F'] = df['var_m_sheet'] * df['lev']
    ols = sm.ols('er ~ -1 + F + F : expected_loss + F : weight_peril', data=df).fit(
        cov_type='cluster', cov_kwds={'groups': df['CUSIP6']}, use_t=True)
    const.append(ols.params[0])
    const_se.append(ols.bse[0])
    el_coef.append(ols.params[1])
    el_coef_se.append(ols.bse[1])
    w_coef.append(ols.params[2])
    w_coef_se.append(ols.bse[2])
    R2 += [ols.rsquared]
    N += [ols.nobs]

    df['F'] = df['lev']
    ols = sm.ols('er ~ -1 + F + F : expected_loss + F : weight_peril', data=df).fit(
        cov_type='cluster', cov_kwds={'groups': df['CUSIP6']}, use_t=True)

    const.append(ols.params[0])
    const_se.append(ols.bse[0])
    el_coef.append(ols.params[1])
    el_coef_se.append(ols.bse[1])
    w_coef.append(ols.params[2])
    w_coef_se.append(ols.bse[2])
    R2 += [ols.rsquared]
    N += [ols.nobs]

    df['F'] = df['er_m']
    ols = sm.ols('er ~ -1 + F + F : expected_loss + F : weight_peril', data=df).fit(
        cov_type='cluster', cov_kwds={'groups': df['CUSIP6']}, use_t=True)
    const.append(ols.params[0])
    const_se.append(ols.bse[0])
    el_coef.append(ols.params[1])
    el_coef_se.append(ols.bse[1])
    w_coef.append(ols.params[2])
    w_coef_se.append(ols.bse[2])
    R2 += [ols.rsquared]
    N += [ols.nobs]

    # %% Print
    data2 = pd.DataFrame([const, const_se, el_coef, el_coef_se, w_coef, w_coef_se])

    T = Table(7, width='auto', justs=['l', 'd', 'd', 'd', 'd', 'd', 'd'],
              caption='Pricing of catastrophe market risk (pooled regressions)',
              label= tname, footer=None)

    T.add_panel(data2, formatting="%.2f", supheader_cols=[1, (2, 4), (5, 7)],
                supheaders=['$\\beta_{i,t}$', '$\\hat{\\beta}_{i,t}$', '$\\phi_0+\\phi_1 el_i + \\phi_2 w_{c,t}$'],
                header_cols=[(1, 1), (2, 2), (3, 3), (4, 4), (5, 5), (6, 6), (7, 7)],
                headers=['$F_t$', '$\\text{Var}_t(R^e_{cat,t+1})\\frac{Size_t}{AUM_t}$', '$\\frac{Size_t}{AUM_t}$',
                         '$E_t(R^e_{cat,t+1})$',
                         '$\\text{Var}_t(R^e_{cat,t+1})\\frac{Size_t}{AUM_t}$', '$\\frac{Size_t}{AUM_t}$',
                         '$E_t(R^e_{cat,t+1})$'],
                regression_specs={'stars': (0.1, 0.05, 0.01), 'dof': N, 'skip_cols': [1]})

    T.add_panel(pd.DataFrame([R2, N]))

    if output == 'paper':
        T.print_table(f'cat_rfs/output/tables/{tname}.tex')
    else:
        T.print_table()

    pass
